package com.ksc.wordcount.driver;

import com.ksc.wordcount.shuffle.ShuffleBlockId;
import com.ksc.wordcount.task.map.MapStatus;
import com.ksc.wordcount.task.TaskContext;
import com.ksc.wordcount.task.TaskStatus;
import com.ksc.wordcount.task.TaskStatusEnum;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.BlockingQueue;

public class TaskManager {

    /**
     * stageId和task队列的映射
     */
    private Map<Integer, BlockingQueue<TaskContext>> stageIdToBlockingQueueMap = new HashMap<>();

    /**
     * stageId和taskId的映射
     */
    private Map<Integer, List<Integer>> stageMap = new HashMap<>();

    /**
     * taskId和task状态的映射
     */
    private Map<Integer, TaskStatus> taskStatusMap = new HashMap<>();

    /**
     * 获取指定stageId的任务队列
     *
     * @param stageId stage的ID
     * @return 任务队列
     */
    public BlockingQueue<TaskContext> getBlockingQueue(int stageId) {
        return stageIdToBlockingQueueMap.get(stageId);
    }

    /**
     * 注册指定stageId的任务队列
     *
     * @param stageId        stage的ID
     * @param blockingQueue  任务队列
     */
    public void registerBlockingQueue(int stageId, BlockingQueue blockingQueue) {
        stageIdToBlockingQueueMap.put(stageId, blockingQueue);
    }

    /**
     * 添加任务上下文到指定stageId的任务队列
     *
     * @param stageId      stage的ID
     * @param taskContext  任务上下文
     */
    public void addTaskContext(int stageId, TaskContext taskContext) {
        // 添加任务到指定的stageId对应的阻塞队列
        BlockingQueue<TaskContext> blockingQueue = getBlockingQueue(stageId);
        blockingQueue.offer(taskContext);

        // 建立stageId和任务Id的映射
        // 如果stageMap没有存在这个stageId,就为其创建一个新的映射
        if (stageMap.get(stageId) == null) {
            stageMap.put(stageId, new ArrayList());
        }
        // 有这个stageId，就把TaskId添加到这个stageId对应的数组集合中
        stageMap.get(stageId).add(taskContext.getTaskId());
    }

    /**
     * 获取指定stageId的任务状态
     *
     * @param stageId  stage的ID
     * @return 任务状态
     */
    public StageStatusEnum getStageTaskStatus(int stageId) {
        //todo 实现 获取指定stage的执行状态，如果该stage下的所有task均执行成功，返回FINISHED
        for (int taskId : stageMap.get(stageId)) {
            if (taskStatusMap.get(taskId) == null) {
                return StageStatusEnum.RUNNING;
            }
            if (taskStatusMap.get(taskId).getTaskStatus() == TaskStatusEnum.FAILED) {
                return StageStatusEnum.FAILED;
            }
            if (taskStatusMap.get(taskId).getTaskStatus() == TaskStatusEnum.RUNNING) {
                return StageStatusEnum.RUNNING;
            }
            if (taskStatusMap.get(taskId).getTaskStatus() == TaskStatusEnum.FINISHED) {
                continue;
            }
        }
        return StageStatusEnum.FINISHED;
    }

    /**
     * 根据reduceId获取指定stageId的ShuffleBlockId数组
     *
     * @param stageId    stage的ID
     * @param reduceId   reduce的ID
     * @return ShuffleBlockId数组
     */
    public ShuffleBlockId[] getStageShuffleIdByReduceId(int stageId, int reduceId) {
        List<ShuffleBlockId> shuffleBlockIds = new ArrayList<>();

        // 遍历指定stageId的所有任务
        for (int taskId : stageMap.get(stageId)) {
            // 获取任务的状态
            TaskStatus taskStatus = taskStatusMap.get(taskId);
            // 将任务状态转换为MapStatus
            MapStatus mapStatus = (MapStatus) taskStatus;

            // 根据reduceId获取对应的ShuffleBlockId
            ShuffleBlockId shuffleBlockId = mapStatus.getShuffleBlockIds()[reduceId];

            // 将ShuffleBlockId添加到列表中
            shuffleBlockIds.add(shuffleBlockId);
        }

        // 将ShuffleBlockId列表转换为数组并返回
        return shuffleBlockIds.toArray(new ShuffleBlockId[shuffleBlockIds.size()]);
    }

    /**
     * 更新任务状态
     *
     * @param taskStatus 任务状态
     */
    public void updateTaskStatus(TaskStatus taskStatus) {
        taskStatusMap.put(taskStatus.getTaskId(), taskStatus);
    }

    private int maxTaskId = 0;

    /**
     *生成任务ID
     *
     * @return 任务ID
     */
    public int generateTaskId() {
        return maxTaskId++;
    }
}